import open3d as o3d
import laspy
import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial import cKDTree
from collections import defaultdict

# ----------------------------
# 参数配置
# ----------------------------
las_path = "/mnt/d/temp_files/part_pointCloud.las"
voxel_size = 0.3
grid_size = 10.0
height_threshold = 1.0
percentile = 5
normal_threshold = 0.95

# 强度范围（道路典型值）
intensity_low = 1000
intensity_high = 42000

# 均匀性参数
search_radius = 1.0      # 局部邻域半径（米）
std_threshold = 3000     # 强度标准差阈值（越小越严格）

guli_search_radius = 1.5  # 孤立点搜索半径（米）

# ----------------------------
# 1-5. 前面的地面提取流程（不变）
# ----------------------------
las = laspy.read(las_path)
points = np.vstack((las.x, las.y, las.z)).T
if hasattr(las, 'intensity'):
    intensity = np.asarray(las.intensity, dtype=np.float64)
else:
    raise ValueError("LAS 文件中没有 'intensity' 字段！")

pcd_orig = o3d.geometry.PointCloud()
pcd_orig.points = o3d.utility.Vector3dVector(points)

if voxel_size is not None:
    pcd = pcd_orig.voxel_down_sample(voxel_size=voxel_size)
    down_points = np.asarray(pcd.points)
    tree_orig = cKDTree(points)
    _, idx = tree_orig.query(down_points, k=1)
    down_intensity = intensity[idx]
else:
    down_points = points
    down_intensity = intensity
    pcd = pcd_orig

print(f"降采样后点数: {down_points.shape[0]}")

# 法线估计
pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamHybrid(radius=1.0, max_nn=30))
pcd.orient_normals_to_align_with_direction(orientation_reference=[0, 0, 1])
normals = np.asarray(pcd.normals)

# 局部高度筛选
x_min, y_min = down_points[:, 0].min(), down_points[:, 1].min()
grid_x = ((down_points[:, 0] - x_min) // grid_size).astype(int)
grid_y = ((down_points[:, 1] - y_min) // grid_size).astype(int)

blocks = defaultdict(list)
for i in range(len(down_points)):
    blocks[(grid_x[i], grid_y[i])].append(i)

height_candidate = np.zeros(len(down_points), dtype=bool)
for key, indices in blocks.items():
    indices = np.array(indices)
    z_vals = down_points[indices, 2]
    local_ground_z = np.percentile(z_vals, percentile)
    mask = down_points[indices, 2] <= (local_ground_z + height_threshold)
    height_candidate[indices[mask]] = True

normal_ok = normals[:, 2] >= normal_threshold
final_ground_mask = height_candidate & normal_ok
ground_points = down_points[final_ground_mask]
ground_intensity = down_intensity[final_ground_mask]

print(f"地面点数: {ground_points.shape[0]}")

# # 添加强度的直方图
# fig, ax = plt.subplots(figsize=(8, 4))
# ax.hist(ground_intensity, bins=100, color='blue', alpha=0.7)
# ax.set_title('Intensity Distribution')
# ax.set_xlabel('Intensity')
# ax.set_ylabel('Frequency')
# plt.show()

if ground_points.shape[0] == 0:
    print("无地面点！")
else:
    # ----------------------------
    # 6. 道路候选：强度在 [1000, 40000]
    # ----------------------------
    road_candidate_mask = (ground_intensity >= intensity_low) & (ground_intensity <= intensity_high)
    road_candidate_points = ground_points[road_candidate_mask]
    road_candidate_intensity = ground_intensity[road_candidate_mask]

    print(f"强度在 [{intensity_low}, {intensity_high}] 的地面点数: {road_candidate_points.shape[0]}")

    if road_candidate_points.shape[0] == 0:
        print("无道路候选点！")
        road_points = np.empty((0, 3))
    else:
        # ----------------------------
        # 7. 计算局部强度标准差（均匀性）
        # ----------------------------
        tree_road = cKDTree(road_candidate_points)
        local_std = np.zeros(road_candidate_points.shape[0])

        for i, pt in enumerate(road_candidate_points):
            # 查找半径内的邻近点
            indices = tree_road.query_ball_point(pt, r=search_radius)
            if len(indices) < 3:  # 至少3个点才能算标准差
                local_std[i] = np.inf
            else:
                neighbor_int = road_candidate_intensity[indices]
                local_std[i] = np.std(neighbor_int)

        # 筛选均匀区域（标准差小）
        uniform_mask = local_std <= std_threshold
        road_points = road_candidate_points[uniform_mask]
        road_intensity = road_candidate_intensity[uniform_mask]

        print(f"道路点数（强度均匀）: {road_points.shape[0]}")

    # ----------------------------
    # 新增：剔除孤立点（邻域点数 < 10）
    # ----------------------------
    if road_points.shape[0] > 0:
        min_neighbors = 10  # 最小邻域点数
        tree_final = cKDTree(road_points)
        # 查询每个点在 search_radius 内的邻居数量（包括自己）
        neighbor_counts = tree_final.query_ball_point(road_points, r=guli_search_radius, return_length=True)
        # 剔除邻居数 < min_neighbors 的点
        dense_mask = np.array(neighbor_counts) >= min_neighbors
        road_points = road_points[dense_mask]
        road_intensity = road_intensity[dense_mask]
        print(f"剔除孤立点后道路点数: {road_points.shape[0]}")
    # ----------------------------

    # 8. 可视化：仅道路点（按强度着色）
    # ----------------------------
    if road_points.shape[0] > 0:
        # 归一化强度用于着色（可选）
        int_norm = (road_intensity - road_intensity.min()) / (road_intensity.max() - road_intensity.min() + 1e-8)
        cmap = plt.get_cmap('plasma')
        road_colors = cmap(int_norm)[:, :3]

        pcd_road = o3d.geometry.PointCloud()
        pcd_road.points = o3d.utility.Vector3dVector(road_points)
        pcd_road.colors = o3d.utility.Vector3dVector(road_colors)

    #     o3d.visualization.draw_geometries(
    #         [pcd_road],
    #         window_name="Extracted Road Points (Intensity + Uniformity)",
    #         width=1200,
    #         height=800
    #     )

        # # ----------------------------
        # # 9. 绘制道路点 XYZ 直方图
        # # ----------------------------
        # fig, axes = plt.subplots(1, 3, figsize=(18, 5))
        # axes[0].hist(road_points[:, 0], bins=80, color='red', alpha=0.7)
        # axes[0].set_title('Road X Distribution')
        # axes[1].hist(road_points[:, 1], bins=80, color='green', alpha=0.7)
        # axes[1].set_title('Road Y Distribution')
        # axes[2].hist(road_points[:, 2], bins=80, color='blue', alpha=0.7)
        # axes[2].set_title('Road Z (Height) Distribution')
        # for ax in axes:
        #     ax.set_ylabel('Frequency')
        # plt.tight_layout()
        # plt.show()

        # ----------------------------
        # 10. 统计信息
        # ----------------------------
        print("\n=== 道路点统计 ===")
        print(f"平均强度: {road_intensity.mean():.1f}")
        print(f"强度标准差（全局）: {road_intensity.std():.1f}")
        print(f"局部强度标准差阈值: {std_threshold}")
        print(f"道路点数: {road_points.shape[0]}")

        # ----------------------------
        # 9. 建筑物识别：通过地面空洞定位
        # ----------------------------
        if ground_points.shape[0] > 0:
            print("\n=== 开始建筑物识别（基于地面空洞）===")
            
            # 参数
            cell_size = 1.5  # 栅格分辨率（米）
            min_building_area = 36  # 最小建筑投影面积（平方米），如 4x4=16
            max_building_area = 500.0  # 最大建筑投影面积
            height_buffer = 0.5  # 地面以上缓冲（米），避免漏掉贴近地面的墙

            # 1. 构建地面覆盖栅格
            x_min, y_min = ground_points[:, 0].min(), ground_points[:, 1].min()
            x_max, y_max = ground_points[:, 0].max(), ground_points[:, 1].max()

            grid_width = int((x_max - x_min) / cell_size) + 1
            grid_height = int((y_max - y_min) / cell_size) + 1
            occupancy = np.zeros((grid_height, grid_width), dtype=bool)

            # 标记有地面点的栅格
            grid_x = ((ground_points[:, 0] - x_min) / cell_size).astype(int)
            grid_y = ((ground_points[:, 1] - y_min) / cell_size).astype(int)
            # 防止越界
            grid_x = np.clip(grid_x, 0, grid_width - 1)
            grid_y = np.clip(grid_y, 0, grid_height - 1)
            occupancy[grid_y, grid_x] = True
            # 可视化occupancy
            # plt.imshow(occupancy)
            # plt.show()


            print(f"地面栅格大小: {grid_width} x {grid_height}")

            # 2. 反转得到空洞（False 表示无地面 → 可能是建筑/树）
            holes = ~occupancy

            # 3. 形态学闭运算：填充小噪声空洞（如车辆、小树）
            from scipy.ndimage import binary_closing, label
            structure = np.ones((1, 1), dtype=bool)  # 3x3 结构元
            holes_clean = binary_closing(holes, structure=structure)

            # 4. 连通域分析
            labeled_holes, num_holes = label(holes_clean)
            print(f"检测到 {num_holes} 个空洞区域")

            building_boxes = []  # 存储 [x_min, y_min, x_max, y_max, base_z]


            # 可视化 occupancy 并绘制所有空洞 bounding box
            fig, ax = plt.subplots(1, 1, figsize=(10, 8))
            # 显示 occupancy（地面为白色，空洞为黑色）
            ax.imshow(occupancy, cmap='gray_r', origin='lower', extent=[x_min, x_max, y_min, y_max])

            # 存储被保留的建筑索引（用于绿色框）
            valid_building_indices = set()

            # 先收集 valid buildings（和你后面逻辑一致）
            for i in range(1, num_holes + 1):
                region_mask = (labeled_holes == i)
                area = np.sum(region_mask) * (cell_size ** 2)
                if min_building_area <= area <= max_building_area:
                    valid_building_indices.add(i)

            # 绘制每个空洞的 bounding box
            for i in range(1, num_holes + 1):
                region_mask = (labeled_holes == i)
                ys, xs = np.where(region_mask)
                if len(xs) == 0:
                    continue
                    
                x1_pix, x2_pix = xs.min(), xs.max()
                y1_pix, y2_pix = ys.min(), ys.max()

                # 转换为世界坐标
                x1_w = x_min + x1_pix * cell_size
                x2_w = x_min + (x2_pix + 1) * cell_size
                y1_w = y_min + y1_pix * cell_size
                y2_w = y_min + (y2_pix + 1) * cell_size

                # 判断是否被保留
                is_valid = i in valid_building_indices
                color = 'green' if is_valid else 'red'
                linewidth = 2 if is_valid else 1

                # 绘制矩形框
                rect = plt.Rectangle((x1_w, y1_w), x2_w - x1_w, y2_w - y1_w,
                                    linewidth=linewidth, edgecolor=color, facecolor='none')
                
                # 根据矩形框的长宽比，判断是否为建筑物
                width = x2_pix - x1_pix + 1
                height = y2_pix - y1_pix + 1

                if width == 0 or height == 0:
                    continue

                aspect_ratio = max(width, height) / min(width, height)
                if aspect_ratio > 1.6:
                    # 长宽比大于 1.6，不是建筑物，直接跳过
                    continue
                ax.add_patch(rect)

            ax.set_title(f'Ground Occupancy with Holes\nGreen: Valid Buildings ({len(valid_building_indices)}), Red: Filtered')
            ax.set_xlabel('X')
            ax.set_ylabel('Y')
            plt.tight_layout()
            plt.show()

            for i in range(1, num_holes + 1):
                region_mask = (labeled_holes == i)
                area = np.sum(region_mask) * (cell_size ** 2)
                
                if area < min_building_area or area > max_building_area:
                    continue

                # 获取区域边界（像素坐标）
                ys, xs = np.where(region_mask)
                x1_pix, x2_pix = xs.min(), xs.max()
                y1_pix, y2_pix = ys.min(), ys.max()

                # 根据矩形框的长宽比，判断是否为建筑物
                width = x2_pix - x1_pix + 1
                height = y2_pix - y1_pix + 1

                if width == 0 or height == 0:
                    continue

                aspect_ratio = max(width, height) / min(width, height)
                if aspect_ratio > 1.6:
                    # 长宽比大于 1.6，不是建筑物，直接跳过
                    continue

                # 转换为世界坐标
                x1 = x_min + x1_pix * cell_size
                x2 = x_min + (x2_pix + 1) * cell_size
                y1 = y_min + y1_pix * cell_size
                y2 = y_min + (y2_pix + 1) * cell_size

                

                # 估算该区域地面高程（取周围地面点的中位数）
                margin = 2.0  # 向外扩展 margin 米取地面点
                around_mask = (
                    (ground_points[:, 0] >= x1 - margin) &
                    (ground_points[:, 0] <= x2 + margin) &
                    (ground_points[:, 1] >= y1 - margin) &
                    (ground_points[:, 1] <= y2 + margin)
                )
                if np.any(around_mask):
                    base_z = np.median(ground_points[around_mask, 2])
                else:
                    # 若周围无地面点，用全局地面中位数
                    base_z = np.median(ground_points[:, 2])

                building_boxes.append([x1, y1, x2, y2, base_z])

            print(f"初步筛选出 {len(building_boxes)} 个建筑候选区域")

            # ----------------------------
            # 10. 从原始 LAS 点云中提取建筑物 3D 点
            # ----------------------------
            if len(building_boxes) > 0:
                all_building_points = []
                all_building_intensities = []

                for idx, (x1, y1, x2, y2, base_z) in enumerate(building_boxes):
                    z_min = base_z - 0.5
                    z_max = base_z + 15.0
                    building_mask = (
                        (points[:, 0] >= x1) & (points[:, 0] <= x2) &
                        (points[:, 1] >= y1) & (points[:, 1] <= y2) &
                        (points[:, 2] >= z_min) & (points[:, 2] <= z_max)
                    )
                    building_pts = points[building_mask]
                    building_ints = intensity[building_mask]  # 提取对应强度
                    if building_pts.shape[0] == 0:
                        continue
                    all_building_points.append(building_pts)
                    all_building_intensities.append(building_ints)
                    
                    if building_pts.shape[0] == 0:
                        continue
                    # 打印建筑的xyz范围
                    print(f"建筑{idx+1}：{building_pts.shape[0]}个点，xyz范围：{building_pts.min(axis=0)} - {building_pts.max(axis=0)}")

                    all_building_points.append(building_pts)


                if all_building_points:
                    final_building_points = np.vstack(all_building_points)
                    final_building_intensities = np.concatenate(all_building_intensities)

                    # 归一化强度并映射颜色
                    int_min, int_max = final_building_intensities.min(), final_building_intensities.max()
                    print(f"建筑物强度范围：{int_min:.1f} - {int_max:.1f}")
                    if int_max > int_min:
                        int_norm = (final_building_intensities - int_min) / (int_max - int_min + 1e-8)
                    else:
                        int_norm = np.zeros_like(final_building_intensities)
                    
                    cmap = plt.get_cmap('hot')  # 也可换 'jet', 'viridis', 'gray' 等
                    building_colors = cmap(int_norm)[:, :3]

                    pcd_buildings = o3d.geometry.PointCloud()
                    pcd_buildings.points = o3d.utility.Vector3dVector(final_building_points)
                    pcd_buildings.colors = o3d.utility.Vector3dVector(building_colors)

                    # ===== 2. 道路点云（蓝色）=====
                    road_colors = np.tile([0.0, 0.0, 1.0], (road_points.shape[0], 1))  # 蓝色
                    pcd_road.colors = o3d.utility.Vector3dVector(road_colors)

                    # ===== 3. 其余点云（灰色）=====
                    is_building = np.zeros(points.shape[0], dtype=bool)
                    for (x1, y1, x2, y2, base_z) in building_boxes:
                        z_min = base_z - 0.5
                        z_max = base_z + 15.0
                        mask = (
                            (points[:, 0] >= x1) & (points[:, 0] <= x2) &
                            (points[:, 1] >= y1) & (points[:, 1] <= y2) &
                            (points[:, 2] >= z_min) & (points[:, 2] <= z_max)
                        )
                        is_building |= mask

                    if road_points.shape[0] > 0:
                        tree_road_orig = cKDTree(road_points)
                        dists, _ = tree_road_orig.query(points, k=1, distance_upper_bound=voxel_size * 2 + 0.1)
                        is_road = np.isfinite(dists)
                    else:
                        is_road = np.zeros(points.shape[0], dtype=bool)

                    other_points = points[~(is_building | is_road)]
                    pcd_other = o3d.geometry.PointCloud()
                    pcd_other.points = o3d.utility.Vector3dVector(other_points)
                    pcd_other.colors = o3d.utility.Vector3dVector(
                        np.full((other_points.shape[0], 3), 0.5)
                    )

                    # ===== 4. 建筑物红色线框（投影框）=====
                    line_sets = []
                    for (x1, y1, x2, y2, base_z) in building_boxes:
                        z_line = base_z + 15.0
                        corners = np.array([
                            [x1, y1, z_line],
                            [x2, y1, z_line],
                            [x2, y2, z_line],
                            [x1, y2, z_line]
                        ])
                        lines = [[0, 1], [1, 2], [2, 3], [3, 0]]
                        colors = [[1.0, 0.0, 0.0] for _ in range(4)]  # 红色

                        line_set = o3d.geometry.LineSet()
                        line_set.points = o3d.utility.Vector3dVector(corners)
                        line_set.lines = o3d.utility.Vector2iVector(lines)
                        line_set.colors = o3d.utility.Vector3dVector(colors)
                        line_sets.append(line_set)

                    # ===== 5. 可视化 =====
                    geometries = [pcd_other, pcd_road, pcd_buildings] + line_sets
                    o3d.visualization.draw_geometries(
                        geometries,
                        window_name="Road (Blue) + Buildings (Intensity Colored) + Boxes (Yellow)",
                        width=1200,
                        height=800
                    )

                    print(f"成功提取 {len(building_boxes)} 栋建筑，共 {final_building_points.shape[0]} 个点")
                else:
                    print("未在原始点云中找到建筑点。")
            else:
                print("无建筑候选区域。")
        else:
            print("无地面点，跳过建筑物识别。")
    else:
        print("未检测到满足均匀性条件的道路点。")